package Algorithm.prefix;

import DataStructure.tree.RestoreBinaryTree;
import DataStructure.tree.TreeNode;

import java.util.ArrayDeque;
import java.util.Deque;
import java.util.HashMap;

/**
 * 437. 路径总和 III https://leetcode.cn/problems/path-sum-iii/
 * 题目简述：统计二叉树中节点值之和为target的路径的数目，这里的路径只能从父节点到子节点
 */
public class PathSum {

    public static void main(String[] args) {
        TreeNode tree = RestoreBinaryTree.createTree("10,5,-3,3,2,null,11,3,-2,null,1");
        new PathSum().pathSum(tree, 8);
    }

    /**
     * 思路：前缀和 + 回溯。使用前序遍历扫描二叉树，并向前缀和集中加入当前前缀和，并分别减去前面所有前缀和求得包含当前节点的所有路径和是否为target。
     * 再递归左子树和右子树。最后扫描完当前节点子树后，将当前前缀和移除。
     */
    int targetSum;
    int count = 0;
    public int pathSum(TreeNode root, int targetSum) {
        this.targetSum = targetSum;
        Deque<Long> preSumDeque = new ArrayDeque<>();
        preSumDeque.add(0L);
        recur(root, preSumDeque);
        return count;
    }

    public void recur(TreeNode root, Deque<Long> preSumDeque) {
        if (root == null) return;
        long sum = preSumDeque.peekLast() + root.val;
        //判断包含当前节点的所有路径和
        for (long num : preSumDeque) {
            if (sum - num == targetSum)
                count++;
        }
        preSumDeque.add(sum);
        recur(root.left, preSumDeque);
        recur(root.right, preSumDeque);
        preSumDeque.removeLast();//回溯移除
    }

    /**
     * 优化：由于只需要统计数量，故可用一个HashMap保存每种前缀和的数量
     */
    public int pathSum2(TreeNode root, int targetSum) {
        this.targetSum = targetSum;
        HashMap<Long, Integer> preSumCount = new HashMap<>();
        preSumCount.put(0L, 1);
        recur2(root, preSumCount, 0);
        return count;
    }
    public void recur2(TreeNode root, HashMap<Long, Integer> preSumCount, long sum) {
        if (root == null) return;
        sum += root.val;
        count += preSumCount.getOrDefault(sum - targetSum, 0);
        preSumCount.put(sum, preSumCount.getOrDefault(sum, 0) + 1);
        recur2(root.left, preSumCount, sum);
        recur2(root.right, preSumCount, sum);
        if (preSumCount.get(sum) == 1) preSumCount.remove(sum);
        else preSumCount.put(sum, preSumCount.get(sum) - 1);
    }
}
